package mqtt311

import (
    "bytes"
    "encoding/binary"
    "fmt"
    "io"
    "unicode/utf8"
    
    "github.com/yangyongzhen/mqtt-broker/internal/protocol/common"
)

// 固定头部
type FixedHeader struct {
    MessageType common.MessageType
    DUP         bool
    QoS         common.QoS
    Retain      bool
    Length      uint32
}

// CONNECT 数据包
type ConnectPacket struct {
    FixedHeader
    ProtocolName    string
    ProtocolVersion byte
    CleanSession    bool
    WillFlag        bool
    WillQoS         common.QoS
    WillRetain      bool
    PasswordFlag    bool
    UsernameFlag    bool
    KeepAlive       uint16
    ClientID        string
    WillTopic       string
    WillMessage     []byte
    Username        string
    Password        []byte
}

func (p *ConnectPacket) Type() common.MessageType {
    return common.CONNECT
}

func (p *ConnectPacket) Encode() ([]byte, error) {
    var buf bytes.Buffer
    
    // 可变头部
    var variableHeader bytes.Buffer
    
    // 协议名称
    writeString(&variableHeader, p.ProtocolName)
    
    // 协议版本
    variableHeader.WriteByte(p.ProtocolVersion)
    
    // 连接标志
    var connectFlags byte
    if p.CleanSession {
        connectFlags |= 0x02
    }
    if p.WillFlag {
        connectFlags |= 0x04
        connectFlags |= byte(p.WillQoS) << 3
        if p.WillRetain {
            connectFlags |= 0x20
        }
    }
    if p.PasswordFlag {
        connectFlags |= 0x40
    }
    if p.UsernameFlag {
        connectFlags |= 0x80
    }
    variableHeader.WriteByte(connectFlags)
    
    // 保活时间
    binary.Write(&variableHeader, binary.BigEndian, p.KeepAlive)
    
    // 载荷
    var payload bytes.Buffer
    writeString(&payload, p.ClientID)
    
    if p.WillFlag {
        writeString(&payload, p.WillTopic)
        writeBytes(&payload, p.WillMessage)
    }
    
    if p.UsernameFlag {
        writeString(&payload, p.Username)
    }
    
    if p.PasswordFlag {
        writeBytes(&payload, p.Password)
    }
    
    // 计算剩余长度
    remainingLength := variableHeader.Len() + payload.Len()
    
    // 固定头部
    fixedHeaderByte := byte(p.MessageType) << 4
    buf.WriteByte(fixedHeaderByte)
    
    // 剩余长度编码
    encodeRemainingLength(&buf, uint32(remainingLength))
    
    // 写入可变头部和载荷
    buf.Write(variableHeader.Bytes())
    buf.Write(payload.Bytes())
    
    return buf.Bytes(), nil
}

func (p *ConnectPacket) Decode(data []byte) error {
    reader := bytes.NewReader(data)
    
    // 读取协议名称
    protocolName, err := readString(reader)
    if err != nil {
        return err
    }
    p.ProtocolName = protocolName
    
    // 读取协议版本
    version, err := reader.ReadByte()
    if err != nil {
        return err
    }
    p.ProtocolVersion = version
    
    // 读取连接标志
    connectFlags, err := reader.ReadByte()
    if err != nil {
        return err
    }
    
    p.CleanSession = (connectFlags & 0x02) != 0
    p.WillFlag = (connectFlags & 0x04) != 0
    p.WillQoS = common.QoS((connectFlags >> 3) & 0x03)
    p.WillRetain = (connectFlags & 0x20) != 0
    p.PasswordFlag = (connectFlags & 0x40) != 0
    p.UsernameFlag = (connectFlags & 0x80) != 0
    
    // 读取保活时间
    err = binary.Read(reader, binary.BigEndian, &p.KeepAlive)
    if err != nil {
        return err
    }
    
    // 读取客户端ID
    clientID, err := readString(reader)
    if err != nil {
        return err
    }
    p.ClientID = clientID
    
    // 读取遗嘱信息
    if p.WillFlag {
        p.WillTopic, err = readString(reader)
        if err != nil {
            return err
        }
        p.WillMessage, err = readBytes(reader)
        if err != nil {
            return err
        }
    }
    
    // 读取用户名
    if p.UsernameFlag {
        p.Username, err = readString(reader)
        if err != nil {
            return err
        }
    }
    
    // 读取密码
    if p.PasswordFlag {
        p.Password, err = readBytes(reader)
        if err != nil {
            return err
        }
    }
    
    return nil
}

// CONNACK 数据包
type ConnackPacket struct {
    FixedHeader
    SessionPresent bool
    ReturnCode     byte
}

func (p *ConnackPacket) Type() common.MessageType {
    return common.CONNACK
}

func (p *ConnackPacket) Encode() ([]byte, error) {
    var buf bytes.Buffer
    
    // 固定头部
    fixedHeaderByte := byte(p.MessageType) << 4
    buf.WriteByte(fixedHeaderByte)
    
    // 剩余长度 (CONNACK 总是 2 字节)
    buf.WriteByte(0x02)
    
    // 连接确认标志
    var ackFlags byte
    if p.SessionPresent {
        ackFlags |= 0x01
    }
    buf.WriteByte(ackFlags)
    
    // 返回码
    buf.WriteByte(p.ReturnCode)
    
    return buf.Bytes(), nil
}

func (p *ConnackPacket) Decode(data []byte) error {
    if len(data) < 2 {
        return fmt.Errorf("CONNACK packet too short")
    }
    
    p.SessionPresent = (data[0] & 0x01) != 0
    p.ReturnCode = data[1]
    
    return nil
}

// PUBLISH 数据包
type PublishPacket struct {
    FixedHeader
    TopicName string
    PacketID  uint16
    Payload   []byte
}

func (p *PublishPacket) Type() common.MessageType {
    return common.PUBLISH
}

func (p *PublishPacket) Encode() ([]byte, error) {
    var buf bytes.Buffer
    var variableHeader bytes.Buffer
    
    // 主题名称
    writeString(&variableHeader, p.TopicName)
    
    // 如果 QoS > 0，包含包标识符
    if p.QoS > 0 {
        binary.Write(&variableHeader, binary.BigEndian, p.PacketID)
    }
    
    // 计算剩余长度
    remainingLength := variableHeader.Len() + len(p.Payload)
    
    // 固定头部
    fixedHeaderByte := byte(p.MessageType) << 4
    if p.DUP {
        fixedHeaderByte |= 0x08
    }
    fixedHeaderByte |= byte(p.QoS) << 1
    if p.Retain {
        fixedHeaderByte |= 0x01
    }
    buf.WriteByte(fixedHeaderByte)
    
    // 剩余长度编码
    encodeRemainingLength(&buf, uint32(remainingLength))
    
    // 写入可变头部和载荷
    buf.Write(variableHeader.Bytes())
    buf.Write(p.Payload)
    
    return buf.Bytes(), nil
}

func (p *PublishPacket) Decode(data []byte) error {
    reader := bytes.NewReader(data)
    
    // 读取主题名称
    topicName, err := readString(reader)
    if err != nil {
        return err
    }
    p.TopicName = topicName
    
    // 如果 QoS > 0，读取包标识符
    if p.QoS > 0 {
        err = binary.Read(reader, binary.BigEndian, &p.PacketID)
        if err != nil {
            return err
        }
    }
    
    // 读取载荷
    payload := make([]byte, reader.Len())
    _, err = reader.Read(payload)
    if err != nil {
        return err
    }
    p.Payload = payload
    
    return nil
}

// SUBSCRIBE 数据包
type SubscribePacket struct {
    FixedHeader
    PacketID      uint16
    Subscriptions []common.Subscription
}

func (p *SubscribePacket) Type() common.MessageType {
    return common.SUBSCRIBE
}

func (p *SubscribePacket) Encode() ([]byte, error) {
    var buf bytes.Buffer
    var variableHeader bytes.Buffer
    var payload bytes.Buffer
    
    // 包标识符
    binary.Write(&variableHeader, binary.BigEndian, p.PacketID)
    
    // 订阅列表
    for _, sub := range p.Subscriptions {
        writeString(&payload, sub.Topic)
        payload.WriteByte(byte(sub.QoS))
    }
    
    // 计算剩余长度
    remainingLength := variableHeader.Len() + payload.Len()
    
    // 固定头部 (SUBSCRIBE 的保留位必须是 0010)
    fixedHeaderByte := byte(p.MessageType)<<4 | 0x02
    buf.WriteByte(fixedHeaderByte)
    
    // 剩余长度编码
    encodeRemainingLength(&buf, uint32(remainingLength))
    
    // 写入可变头部和载荷
    buf.Write(variableHeader.Bytes())
    buf.Write(payload.Bytes())
    
    return buf.Bytes(), nil
}

func (p *SubscribePacket) Decode(data []byte) error {
    reader := bytes.NewReader(data)
    
    // 读取包标识符
    err := binary.Read(reader, binary.BigEndian, &p.PacketID)
    if err != nil {
        return err
    }
    
    // 读取订阅列表
    p.Subscriptions = nil
    for reader.Len() > 0 {
        topic, err := readString(reader)
        if err != nil {
            return err
        }
        
        qosByte, err := reader.ReadByte()
        if err != nil {
            return err
        }
        
        p.Subscriptions = append(p.Subscriptions, common.Subscription{
            Topic: topic,
            QoS:   common.QoS(qosByte),
        })
    }
    
    return nil
}

// SUBACK 数据包
type SubackPacket struct {
    FixedHeader
    PacketID    uint16
    ReturnCodes []byte
}

func (p *SubackPacket) Type() common.MessageType {
    return common.SUBACK
}

func (p *SubackPacket) Encode() ([]byte, error) {
    var buf bytes.Buffer
    
    // 固定头部
    fixedHeaderByte := byte(p.MessageType) << 4
    buf.WriteByte(fixedHeaderByte)
    
    // 剩余长度
    remainingLength := 2 + len(p.ReturnCodes)
    encodeRemainingLength(&buf, uint32(remainingLength))
    
    // 包标识符
    binary.Write(&buf, binary.BigEndian, p.PacketID)
    
    // 返回码
    buf.Write(p.ReturnCodes)
    
    return buf.Bytes(), nil
}

func (p *SubackPacket) Decode(data []byte) error {
    if len(data) < 2 {
        return fmt.Errorf("suback packet too short")
    }
    
    // 解析固定头部
    p.PacketID = binary.BigEndian.Uint16(data[0:2])
    
    // 解析返回码
    p.ReturnCodes = make([]byte, len(data)-2)
    copy(p.ReturnCodes, data[2:])
    
    return nil
}

// 工具函数
func writeString(w io.Writer, s string) error {
    if !utf8.ValidString(s) {
        return fmt.Errorf("invalid UTF-8 string")
    }
    length := uint16(len(s))
    binary.Write(w, binary.BigEndian, length)
    _, err := w.Write([]byte(s))
    return err
}

func writeBytes(w io.Writer, b []byte) error {
    length := uint16(len(b))
    binary.Write(w, binary.BigEndian, length)
    _, err := w.Write(b)
    return err
}

func readString(r io.Reader) (string, error) {
    var length uint16
    err := binary.Read(r, binary.BigEndian, &length)
    if err != nil {
        return "", err
    }
    
    data := make([]byte, length)
    _, err = io.ReadFull(r, data)
    if err != nil {
        return "", err
    }
    
    s := string(data)
    if !utf8.ValidString(s) {
        return "", fmt.Errorf("invalid UTF-8 string")
    }
    
    return s, nil
}

func readBytes(r io.Reader) ([]byte, error) {
    var length uint16
    err := binary.Read(r, binary.BigEndian, &length)
    if err != nil {
        return nil, err
    }
    
    data := make([]byte, length)
    _, err = io.ReadFull(r, data)
    return data, err
}

func encodeRemainingLength(w io.Writer, length uint32) error {
    for {
        encodedByte := byte(length % 128)
        length = length / 128
        if length > 0 {
            encodedByte |= 128
        }
        w.Write([]byte{encodedByte})
        if length == 0 {
            break
        }
    }
    return nil
}

func decodeRemainingLength(r io.Reader) (uint32, error) {
    var length uint32
    var multiplier uint32 = 1
    
    for {
        encodedByte := make([]byte, 1)
        _, err := r.Read(encodedByte)
        if err != nil {
            return 0, err
        }
        
        length += uint32(encodedByte[0]&127) * multiplier
        if (encodedByte[0] & 128) == 0 {
            break
        }
        multiplier *= 128
        if multiplier > 128*128*128 {
            return 0, fmt.Errorf("remaining length exceeds maximum")
        }
    }
    
    return length, nil
}
